classdef abilityDistFixed < handle
    properties
        integration_rule_g 
        
        mass % used for normalization in case of truncation
    end
    
    properties (SetAccess = protected)
        num_nodes_f  % integration nodes are pretty much legacy
                     % code - no integration only done when
                     % computing total mass
        num_nodes_g 
        num_nodes_bins 

        nodes_f 
        % pre-computed nodes for integration over :math:`f_i(w)`
        weights_f 
        % pre-computed weights for integration over :math:`f_i(w)`
        nodes_g 
        % pre-computed nodes for integration over :math:`g(\theta_M,\theta_R,\theta_C)`
        weights_g 
        % pre-computed weights for integration over :math:`g(\theta_M,\theta_R,\theta_C)`
        nodes_bins 
        % pre-computed nodes for integration within bins of :math:`f_i(w)`
        weights_bins 
        % pre-computed weights for integration within bins of :math:`f_i(w)`
    end

    methods
        
        function self = abilityDistFixed(varargin)
            p = inputParser;
            p.KeepUnmatched = true;
            
            % functions to check inputs
            valid_num_nodes = @(x) isnumeric(x) && x==floor(x) && ...
                (x>0);
            
            % input requirements
            addParameter(p, 'num_nodes_g', 15, valid_num_nodes)
            addParameter(p, 'num_nodes_f', 10, valid_num_nodes)
            addParameter(p, 'num_nodes_bins', 10, valid_num_nodes)
            addParameter(p, 'integration_rule_g','gauss-legendre')
            addParameter(p, 'mass',1, @(x) x > 0)
            
            parse(p, varargin{:});
            
            self.num_nodes_g = p.Results.num_nodes_g;
            self.num_nodes_f = p.Results.num_nodes_f;
            self.num_nodes_bins = p.Results.num_nodes_bins;

            self.integration_rule_g = p.Results.integration_rule_g;
            self.mass = p.Results.mass;
            
            % pre-compute Gauss-Legendre nodes and weights
            [self.nodes_f, self.weights_f] = tools.Integration.lgwt(self.num_nodes_f,-1,1);
            [self.nodes_g, self.weights_g] = tools.Integration.lgwt(self.num_nodes_g,-1,1);
            [self.nodes_bins, self.weights_bins] = tools.Integration.lgwt(self.num_nodes_bins,-1,1);
        end

        function set_num_nodes(self, num_nodes_f, num_nodes_g, num_nodes_bins)
        % Set number of nodes properties
        %
        % Args:
        %     num_nodes_f (int): number of nodes to use for
        %         Gauss-Legendre integration over :math:`f()`
        %     num_nodes_g (int): number of nodes to use for
        %         Gauss-Legendre integration over :math:`g()`
        %     num_nodes_bins (double): number of nodes to use for
        %         Gauss-Legendre integration over bins of :math:`f()`
            
            p = inputParser;
            p.addRequired('num_nodes_f', @self.valid_num_nodes);
            p.addRequired('num_nodes_g', @self.valid_num_nodes);
            p.addRequired('num_nodes_bins', @self.valid_num_nodes);
            
            p.parse(num_nodes_f, num_nodes_g, num_nodes_bins);
            
            self.num_nodes_f = p.Results.num_nodes_f;
            self.num_nodes_g = p.Results.num_nodes_g;
            self.num_nodes_bins = p.Results.num_nodes_bins;
            
            % pre-compute Gauss-Legendre nodes and weights
            [self.nodes_f, self.weights_f] = tools.Integration.lgwt(self.num_nodes_f,-1,1);
            [self.nodes_g, self.weights_g] = tools.Integration.lgwt(self.num_nodes_g,-1,1);
            [self.nodes_bins, self.weights_bins] = tools.Integration.lgwt(self.num_nodes_bins,-1,1);
        end        
        
        function out = total_mass(self,Y_M, Y_R, Y_C, w_lb, w_ub, par)
        % Compute total mass by integrating over :math:`f(w)` from
        % :code:`w_lb` to :code:`w_ub`
            integrand = @(w) self.f(w,Y_M,Y_R,Y_C, par);
            out = tools.Integration.integral_legendre(integrand, w_lb, w_ub, self.num_nodes_f,self.nodes_f,self.weights_f);
        end       

        function out = f(self, w, Y_M, Y_R, Y_C, par)
        % Wage density :math:`f(w)`.
            out = self.f_M(w, Y_M, Y_R, Y_C, par) + ...
                  self.f_R(w, Y_M, Y_R, Y_C, par) + ...
                  self.f_C( w, Y_M, Y_R, Y_C, par);
        end
        
        function out = f_M(self, w_all, Y_M, Y_R, Y_C, par)
        % Wage density :math:`f_M(w)`.
            out = par.shareM .* 1/Y_M.*plognpdf(w_all./Y_M, par.aM, par.muM, par.sigmaM)./self.mass;
        end
        
        function out = f_R(self, w_all, Y_M, Y_R, Y_C, par)
        % Wage density :math:`f_R(w)`.
            out = par.shareR .* 1/Y_R.*plognpdf(w_all./Y_R, par.aR, par.muR, par.sigmaR)./self.mass;
        end
        
        function out = f_C(self, w_all, Y_M, Y_R, Y_C, par)
        % Wage density :math:`f_C(w)`.
            out = par.shareC .* 1/Y_C.*plognpdf(w_all./Y_C, par.aC, par.muC, par.sigmaC)./self.mass;
        end

         
        %% integration nodes and weights        
        
        function [nodes, weights] = integration_points_weights(self, num_nodes_g)
        % Compute integration points (x,y) and weights w for
        % integration over :math`g()`. Type of nodes and weights
        % depends on property :code:`integration_rule_g`.
        %
        % Arg:
        %     num_nodes_g (int): number of nodes to use for
        %         Gauss-Legendre integration over :math:`g()`
            
            switch self.integration_rule_g
              case 'gauss-legendre'
                if num_nodes_g == self.num_nodes_g
                    % use pre-computed nodes and weights if
                    % number of nodes is as initialized 
                    nodes = self.nodes_g;
                    weights = self.weights_g;
                else
                    [nodes, weights] = tools.Integration.lgwt(num_nodes_g,-1,1);
                end
                
               nodes = nodes';
                
            end
        end
         
    end
end
